package com.chen.scheduler;

import com.alibaba.fastjson.JSON;
import com.chen.pipeline.JDItemPipeline;
import com.chen.util.DbUtil;
import org.apache.commons.codec.digest.DigestUtils;
import org.apache.commons.lang3.StringUtils;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.scheduler.RedisScheduler;

import java.nio.charset.StandardCharsets;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.LinkedBlockingDeque;


public class RedisJDScheduler extends RedisScheduler {
    private static final String ZSET_PREFIX = "zset_";

    private static final String QUEUE_PREFIX = "queue_";

    private static final String NO_PRIORITY_SUFFIX = "_zore";

    private static final String PLUS_PRIORITY_SUFFIX = "_plus";

    private static final String MINUS_PRIORITY_SUFFIX = "_minus";

    /**
     * 本地存储要爬的页数的队列
     */
    private final Queue<String> queue = new LinkedBlockingDeque<>();
    /**
     * 本地要爬取的队列的request信息
     */

    private final HashMap<String, String> map = new HashMap<>();

    public RedisJDScheduler(String host) {
        super(host);
    }

    public RedisJDScheduler(JedisPool pool) {
        super(pool);
    }

    public RedisJDScheduler(JedisPool pool, Properties properties) {
        super(pool);
    }


    @Override
    protected void pushWhenNoDuplicate(Request request, Task task) {
        Jedis jedis = pool.getResource();
        try {
            if (request.getPriority() == 5) {
                queue.offer(request.getUrl());
                setQueueExtrasInItem(request, task);
                return;
            } else if (request.getPriority() > 0)
                jedis.zadd(getZsetPlusPriorityKey(task), request.getPriority(), request.getUrl());
            else if (request.getPriority() < 0)
                jedis.zadd(getZsetMinusPriorityKey(task), request.getPriority(), request.getUrl());
            else {
                DbUtil.brandAdd(request.getUrl(), JSON.toJSONString(request));
                return;
            }

            setExtrasInItem(jedis, request, task);
        } catch (SQLException throwables) {
            throwables.printStackTrace();
        } finally {
            jedis.close();
        }
    }

    @Override
    public synchronized Request poll(Task task) {
        Jedis jedis = pool.getResource();
        try {
            String url;
            url = queue.poll();
            if (!StringUtils.isBlank(url)) {
                return getQueueExtrasInItem(url, task);
            }
            Map.Entry<String, String> brandUrl = DbUtil.brandPoll();
            if (brandUrl != null) {
                return JSON.parseObject(brandUrl.getValue(), Request.class);
            }
            url = getRequest(jedis, task);
            if (StringUtils.isBlank(url)) {
                return null;
            }
            if (getQueueExtrasInItem(url, task) != null) {
                return getQueueExtrasInItem(url, task);
            }
            return getExtrasInItem(jedis, url, task);
        } catch (SQLException throwables) {
            throwables.printStackTrace();
        } finally {
            jedis.close();
        }
        return null;
    }

    private String getRequest(Jedis jedis, Task task) {
        String url;
        Set<String> urls = jedis.zrevrange(getZsetPlusPriorityKey(task), 0, 0);
        if (urls.isEmpty()) {
            url = jedis.lpop(getQueueNoPriorityKey(task));
            if (StringUtils.isBlank(url)) {
                urls = jedis.zrevrange(getZsetMinusPriorityKey(task), 0, 0);
                if (!urls.isEmpty()) {
                    url = urls.toArray(new String[0])[0];
                    jedis.zrem(getZsetMinusPriorityKey(task), url);
                    JDItemPipeline.updateType();
                }
            }
        } else {
            url = urls.toArray(new String[0])[0];
            jedis.zrem(getZsetPlusPriorityKey(task), url);
        }
        return url;
    }

    @Override
    public void resetDuplicateCheck(Task task) {
        Jedis jedis = pool.getResource();
        try {
            jedis.del(getSetKey(task));
        } finally {
            jedis.close();
        }
    }

    @Override
    public boolean isDuplicate(Request request, Task task) {
        Jedis jedis = pool.getResource();
        try {
            if (request.getExtra("retry") != null && request.getExtra("retry").equals("true")) {
                request.putExtra("retry", "false");
                return false;
            }
            if (request.getPriority() == 5) {
                return false;
            }
            if (request.getPriority() == 0) {
                if (DbUtil.isDuplicate(request.getUrl()))
                    return true;
                else
                    return false;
            }
            return jedis.sadd(getSetKey(task), request.getUrl()) == 0;
        } catch (SQLException throwables) {
            throwables.printStackTrace();
        } finally {
            jedis.close();
        }
        return true;
    }

    private String getZsetPlusPriorityKey(Task task) {
        return ZSET_PREFIX + task.getUUID() + PLUS_PRIORITY_SUFFIX;
    }

    private String getQueueNoPriorityKey(Task task) {
        return QUEUE_PREFIX + task.getUUID() + NO_PRIORITY_SUFFIX;
    }

    private String getZsetMinusPriorityKey(Task task) {
        return ZSET_PREFIX + task.getUUID() + MINUS_PRIORITY_SUFFIX;
    }

    private void setExtrasInItem(Jedis jedis, Request request, Task task) {
        if (request.getExtras() != null) {
            String field = DigestUtils.shaHex(request.getUrl());
            String value = JSON.toJSONString(request);
            jedis.hset(getItemKey(task), field, value);
        }
    }

    private Request getExtrasInItem(Jedis jedis, String url, Task task) {
        String key = getItemKey(task);
        String field = DigestUtils.shaHex(url);
        byte[] bytes = jedis.hget(key.getBytes(), field.getBytes());
        if (bytes != null)
            return JSON.parseObject(new String(bytes, StandardCharsets.UTF_8), Request.class);
        return new Request(url);
    }

    private void setQueueExtrasInItem(Request request, Task task) {
        if (request.getExtras() != null) {
            String field = DigestUtils.shaHex(request.getUrl());
            String value = JSON.toJSONString(request);
            map.put(field, value);
        }
    }

    private Request getQueueExtrasInItem(String url, Task task) {
        String field = DigestUtils.shaHex(url);
        String s = map.get(field);
        map.remove(field);
        if (s != null)
            return JSON.parseObject(s, Request.class);
        return null;
    }
}
